#!/usr/bin/env python
# encoding: utf-8
'''
@author: zhaowenpeng
@contact: winston@peipeiyun.com
@software: garner
@file: __init__.py.py
@time: 2020/7/17 3:43 下午
@desc:
'''


from playhouse.pool import PooledMySQLDatabase
from peewee import OperationalError
from peewee import __exception_wrapper__
class RetryOperationalError(object):
    def execute_sql(self, sql, params=None, commit=True):
        try:
            cursor = super(RetryOperationalError, self).execute_sql(
                sql, params, commit)
        except OperationalError:
            if not self.is_closed():
                self.close()
            with __exception_wrapper__:
                cursor = self.cursor()
                cursor.execute(sql, params or ())
                if commit and not self.in_transaction():
                    self.commit()
        return cursor
class RetryMySQLDatabase(RetryOperationalError, PooledMySQLDatabase):
    pass